#include <Windows.h>
#include <shellapi.h>
#pragma comment(lib, "Gdi32.lib")
#pragma comment(lib, "user32.lib")
#pragma comment(lib, "Shell32.lib")

#include <string>
#include <vector>
#include <fstream>
#include <iostream>

struct ICONDIRENTRY {
  UCHAR nWidth;
  UCHAR nHeight;
  UCHAR nNumColorsInPalette; // 0 if no palette
  UCHAR nReserved; // should be 0
  WORD nNumColorPlanes; // 0 or 1
  WORD nBitsPerPixel;
  ULONG nDataLength; // length in bytes
  ULONG nOffset; // offset of BMP or PNG data from beginning of file
};

bool SaveIcon(HICON hIcon, const std::string& path, int nColorBits = 32) {
  if (offsetof(ICONDIRENTRY, nOffset) != 12)
    return false;

  HDC dc = CreateCompatibleDC(NULL);

  std::ofstream file(path, std::ios::binary);
  if (!file.is_open())
    return false;

  // Write header:
  char icoHeader[6] = { 0, 0, 1, 0, 1, 0 }; // ICO file with 1 image
  file.write(icoHeader, sizeof(icoHeader));

  // Get information about icon:
  ICONINFO iconInfo;
  GetIconInfo(hIcon, &iconInfo);
  BITMAPINFO bmInfo = { 0 };
  bmInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
  bmInfo.bmiHeader.biBitCount = 0;    // don't get the color table     
  if (!GetDIBits(dc, iconInfo.hbmColor, 0, 0, NULL, &bmInfo, DIB_RGB_COLORS)) {
    return false;
  }

  // Allocate size of bitmap info header plus space for color table:
  int nBmInfoSize = sizeof(BITMAPINFOHEADER);
  if (nColorBits < 24) {
    nBmInfoSize += sizeof(RGBQUAD) * (int)(1 << nColorBits);
  }

  std::vector<UCHAR> bitmapInfo;
  bitmapInfo.resize(nBmInfoSize);
  BITMAPINFO* pBmInfo = (BITMAPINFO *)bitmapInfo.data();
  memcpy(pBmInfo, &bmInfo, sizeof(BITMAPINFOHEADER));

  // Get bitmap data:
  if (!bmInfo.bmiHeader.biSizeImage)
    return false;
  std::vector<UCHAR> bits;
  bits.resize(bmInfo.bmiHeader.biSizeImage);
  pBmInfo->bmiHeader.biBitCount = nColorBits;
  pBmInfo->bmiHeader.biCompression = BI_RGB;
  if (!GetDIBits(dc, iconInfo.hbmColor, 0, bmInfo.bmiHeader.biHeight, bits.data(), pBmInfo, DIB_RGB_COLORS)) {
    return false;
  }

  // Get mask data:
  BITMAPINFO maskInfo = { 0 };
  maskInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
  maskInfo.bmiHeader.biBitCount = 0;  // don't get the color table     
  if (!GetDIBits(dc, iconInfo.hbmMask, 0, 0, NULL, &maskInfo, DIB_RGB_COLORS) || maskInfo.bmiHeader.biBitCount != 1)
    return false;

  std::vector<UCHAR> maskBits;
  maskBits.resize(maskInfo.bmiHeader.biSizeImage);
  std::vector<UCHAR> maskInfoBytes;
  maskInfoBytes.resize(sizeof(BITMAPINFO) + 2 * sizeof(RGBQUAD));
  BITMAPINFO* pMaskInfo = (BITMAPINFO*)maskInfoBytes.data();
  memcpy(pMaskInfo, &maskInfo, sizeof(maskInfo));
  if (!GetDIBits(dc, iconInfo.hbmMask, 0, maskInfo.bmiHeader.biHeight, maskBits.data(), pMaskInfo, DIB_RGB_COLORS)) {
    return false;
  }

  // Write directory entry:
  ICONDIRENTRY dir;
  dir.nWidth = (UCHAR)pBmInfo->bmiHeader.biWidth;
  dir.nHeight = (UCHAR)pBmInfo->bmiHeader.biHeight;
  dir.nNumColorsInPalette = (nColorBits == 4 ? 16 : 0);
  dir.nReserved = 0;
  dir.nNumColorPlanes = 0;
  dir.nBitsPerPixel = pBmInfo->bmiHeader.biBitCount;
  dir.nDataLength = pBmInfo->bmiHeader.biSizeImage + pMaskInfo->bmiHeader.biSizeImage + nBmInfoSize;
  dir.nOffset = sizeof(dir) + sizeof(icoHeader);
  file.write((const char*)&dir, sizeof(dir));

  // Write DIB header (including color table):
  int nBitsSize = pBmInfo->bmiHeader.biSizeImage;
  pBmInfo->bmiHeader.biHeight *= 2; // because the header is for both image and mask
  pBmInfo->bmiHeader.biCompression = 0;
  pBmInfo->bmiHeader.biSizeImage += pMaskInfo->bmiHeader.biSizeImage; // because the header is for both image and mask
  file.write((const char*)&pBmInfo->bmiHeader, nBmInfoSize);

  // Write image data:
  file.write((const char*)(bits.data()), nBitsSize);

  // Write mask data:
  file.write((const char*)(maskBits.data()), pMaskInfo->bmiHeader.biSizeImage);

  file.close();

  DeleteObject(iconInfo.hbmColor);
  DeleteObject(iconInfo.hbmMask);

  DeleteDC(dc);

  return true;
}

int main(int argc, char **argv) {
  if (argc != 2) {
    std::cout << "give a exe or dll file." << std::endl;
    return -1;
  }

  std::string strIcoPath = argv[1];
  std::string strLargePath = strIcoPath + ".big.";
  std::string strSmallPath = strIcoPath + ".small.";
  std::string tmp;
  
  UINT ok_count = 0;
  UINT count = ExtractIconExA(strIcoPath.c_str(), -1, NULL, NULL, 0);
  for (UINT i = 0; i < count; ++i) {
    HICON hIconLarge = NULL, hIconSmall = NULL;
    ExtractIconExA(strIcoPath.c_str(), i, &hIconLarge, NULL, 1);
    ExtractIconExA(strIcoPath.c_str(), i, NULL, &hIconSmall, 1);
    tmp = std::to_string(i) + ".ico";
    if (hIconLarge != NULL) {
      auto ok = SaveIcon(hIconLarge, strLargePath + tmp);
		  DestroyIcon(hIconLarge);
      if (ok)
        ++ok_count;
    }
    if (hIconSmall != NULL) {
      auto ok = SaveIcon(hIconSmall, strSmallPath + tmp);
		  DestroyIcon(hIconSmall);
      if (ok)
        ++ok_count;
    }
  }

  std::cout << count << " found,";
  std::cout << ok_count << " saved." << std::endl;

  return 0;
}